#include <iostream>
#include <cstring>
#include "PriorityQueue.h"
#include "Stack.h"

#define CHARNUM 128
#define K 5

using namespace std;

struct FREQ{
    char chr;
    int count;

    FREQ():chr(0),count(0)
    {}
};

struct HTNode {
    char  chr;
    int   weight;
    HTNode** child;

    HTNode():chr(0),weight(0),child(new HTNode*[K])
	{
		for (int i = 0; i < K; i++) child[i] = nullptr;
	}
    HTNode(char chr, int weight):chr(chr),weight(weight),child(new HTNode*[K])
    {
		for (int i = 0; i < K; i++) child[i] = nullptr;
	}
};


FREQ* GetFrequence(const char* s, int* count) {
    char ch;
    short charTypeCount = 0;
    int AscII[CHARNUM] = {0};

    while((ch = *s++) != '\0') {
        AscII[ch] += 1;
    }

    for(int i = 0; i < CHARNUM; i++) {
        if(AscII[i] != 0) charTypeCount++;
    }
    *count = charTypeCount;

    FREQ* p = new FREQ[charTypeCount];

    int j = 0;
    for(int i = 0; i < CHARNUM; i++) {
        if(AscII[i] != 0) {
            p[j].chr = i;
            p[j].count = AscII[i];
            j++;
        }
    }

    return p;
}


HTNode* BuildTree(FREQ* freq, int charTypeCount) {
    int k0  = (charTypeCount - 1) % (K - 1);
    int nodesToComplete = K - 1 - k0;
    PriorityQueue<HTNode*> PQ;

    //Insert all node into the priority queue
    for(int i = 0; i < charTypeCount; i++) {
        PQ.insert(new HTNode(freq[i].chr, freq[i].count), -freq[i].count);
    }

    //Complete the nodes
    for(int i = 0; i < nodesToComplete; i++) {
        PQ.insert(new HTNode('\0', 0), 0);
    }

    //create tree from priority queue
    while(1 != PQ.length()) {
        HTNode* p = new HTNode;
        HTNode* temp;
        int weight = 0;
        for(int i = 0; i < K; i++) {
            temp = PQ.pullPeek();
            weight += temp->weight;
            p->child[i] = temp;
        }
		p->weight = weight;
        PQ.insert(p, -weight);
    }

    return PQ.pullPeek();
}

char* SerializeStack(Stack<char>& S) {
	int count = 0;
	StackNode<char>* p = S.head->next;
	//The reverse order
	while (p != nullptr) {
		count++;
		p = p->next;
	}

	char* code = new char[count+1];
	count = 0;
	p = S.head->next;
	while (p != nullptr) {
		code[count] = p->data;
		count++;
		p = p->next;
	}
	code[count] = '\0';
	_strrev(code);
	return code;
}

void GenerateCode(HTNode* kHTree, Stack<char>& S, char** codeTable) {
	//stop when the node is a leaf
	if (kHTree->child[0] == nullptr) {
		if (kHTree->chr != '\0') {
			//generate code from stack
			char* code = SerializeStack(S);
			codeTable[kHTree->chr] = code;
		}
		return;
	}

	for (int i = 0; i < K; i++) {
		S.push(i + '0');
		GenerateCode(kHTree->child[i], S, codeTable);
		S.pop();
	}
}

void DisplayCodeTable(char** codeTable) {
	for (int i = 0; i < 128; i++) {
		if (codeTable[i] != nullptr) {
			cout << (char)i << ": " << codeTable[i] << endl;
		}
	}
}

void Encode(char* s, char** codeTable) {
	char ch;
	while ((ch = *s++) != '\0') {
		cout << codeTable[ch];
	}
}

void Decode(char* s, HTNode* kHTree) {
	char ch;
	HTNode* curNode = kHTree;

	while ((ch = *s) != '\0') {
		if (!curNode->child[0]) {
			cout << curNode->chr;
			curNode = kHTree;
		}
		else {
			curNode = curNode->child[ch - '0'];
			s++;
		}
	}
	if (!curNode->child[0]) {
		cout << curNode->chr;
	}
}

int main() {
    char s[] = "aaaaabbbcccccdddefgggggggggghhhhiiiiiiiiiiiijkkllllmmnooooppppp";
    int charTypeCount;
    FREQ* F = GetFrequence(s, &charTypeCount);
    HTNode* H = BuildTree(F, charTypeCount);

	Stack<char> S;
	char** codeTable = new char*[128];
	for (int i = 0; i < 128; i++) codeTable[i] = nullptr;

	GenerateCode(H, S, codeTable);

	DisplayCodeTable(codeTable);

	char toEncode[100];
	cout << "Input string to encode: ";
	cin >> toEncode;
	Encode(toEncode, codeTable);

	cout << endl;

	char toDecode[100];
	cout << "Input code to decode: ";
	cin >> toDecode;
	Decode(toDecode, H);

	cout << endl;

	return 0;
}
